// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <kllm/core/kv_override.h>
#include <cstring>
#include <unordered_set>
#include <unordered_map>
#include <turbo/log/logging.h>

namespace kllm {


    bool string_parse_kv_override(const char * data, std::vector<llama_model_kv_override> & overrides) {
        const char * sep = strchr(data, '=');
        if (sep == nullptr || sep - data >= 128) {
            LOG(ERROR)<<__func__<<": malformed KV override '"<<data<<"'";
            return false;
        }
        llama_model_kv_override kvo;
        std::strncpy(kvo.key, data, sep - data);
        kvo.key[sep - data] = 0;
        sep++;
        if (strncmp(sep, "int:", 4) == 0) {
            sep += 4;
            kvo.tag = LLAMA_KV_OVERRIDE_TYPE_INT;
            kvo.val_i64 = std::atol(sep);
        } else if (strncmp(sep, "float:", 6) == 0) {
            sep += 6;
            kvo.tag = LLAMA_KV_OVERRIDE_TYPE_FLOAT;
            kvo.val_f64 = std::atof(sep);
        } else if (strncmp(sep, "bool:", 5) == 0) {
            sep += 5;
            kvo.tag = LLAMA_KV_OVERRIDE_TYPE_BOOL;
            if (std::strcmp(sep, "true") == 0) {
                kvo.val_bool = true;
            } else if (std::strcmp(sep, "false") == 0) {
                kvo.val_bool = false;
            } else {
                LOG(ERROR)<<__func__<<": invalid boolean value for KV override '"<<data<<"'";
                return false;
            }
        } else if (strncmp(sep, "str:", 4) == 0) {
            sep += 4;
            kvo.tag = LLAMA_KV_OVERRIDE_TYPE_STR;
            if (strlen(sep) > 127) {
                LOG(WARNING)<<__func__<<": malformed KV override '"<<data<<"', value cannot exceed 127 chars";
                return false;
            }
            strncpy(kvo.val_str, sep, 127);
            kvo.val_str[127] = '\0';
        } else {
            LOG(WARNING)<<__func__<<": invalid type for KV override '"<<data<<"'";
            return false;
        }
        overrides.emplace_back(std::move(kvo));
        return true;
    }


    void common_kv_cache_dump_view(const llama_kv_cache_view & view, int row_size) {
        static const char slot_chars[] = ".123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz+";

        printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, largest empty slot=%d @ %d",
               view.n_cells, view.n_seq_max, view.used_cells, view.token_count, view.max_contiguous, view.max_contiguous_idx);

        llama_kv_cache_view_cell * c_curr = view.cells;
        llama_seq_id * cs_curr = view.cells_sequences;

        for (int i = 0; i < view.n_cells; i++, c_curr++, cs_curr += view.n_seq_max) {
            if (i % row_size == 0) {
                printf("\n%5d: ", i);
            }
            int seq_count = 0;
            for (int j = 0; j < view.n_seq_max; j++) {
                if (cs_curr[j] >= 0) { seq_count++; }
            }
            putchar(slot_chars[std::min(sizeof(slot_chars) - 2, size_t(seq_count))]);
        }

        printf("\n=== Done dumping\n");
    }

    void common_kv_cache_dump_view_seqs(const llama_kv_cache_view & view, int row_size) {
        static const char slot_chars[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";

        printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, largest empty slot=%d @ %d\n",
               view.n_cells, view.n_seq_max, view.used_cells, view.token_count, view.max_contiguous, view.max_contiguous_idx);

        std::unordered_map<llama_seq_id, size_t> seqs;
        llama_kv_cache_view_cell * c_curr = view.cells;
        llama_seq_id * cs_curr = view.cells_sequences;

        for (int i = 0; i < view.n_cells; i++, c_curr++, cs_curr += view.n_seq_max) {
            for (int j = 0; j < view.n_seq_max; j++) {
                if (cs_curr[j] < 0) { continue; }
                if (seqs.find(cs_curr[j]) == seqs.end()) {
                    if (seqs.size() + 1 >= sizeof(slot_chars)) { break; }
                    const size_t sz = seqs.size();
                    seqs[cs_curr[j]] = sz;
                }
            }
            if (seqs.size() + 1 >= sizeof(slot_chars)) { break; }
        }

        printf("=== Sequence legend: ");
        for (const auto & it : seqs) {
            printf("%zu=%d, ", it.second, it.first);
        }
        printf("'+'=other sequence ids");

        c_curr = view.cells;
        cs_curr = view.cells_sequences;
        for (int i = 0; i < view.n_cells; i++, c_curr++, cs_curr += view.n_seq_max) {
            if (i % row_size == 0) {
                printf("\n%5d: ", i);
            }
            for (int j = 0; j < view.n_seq_max; j++) {
                if (cs_curr[j] >= 0) {
                    const auto & it = seqs.find(cs_curr[j]);
                    putchar(it != seqs.end() ? int(slot_chars[it->second]) : '+');
                } else {
                    putchar('.');
                }
            }
            putchar(' ');
        }

        printf("\n=== Done dumping\n");
    }

}